import logging
import os
import sys

import hydra
import numpy as np
from omegaconf import OmegaConf
from tensorboardX import SummaryWriter
from tqdm import trange

from model.data_processing.jax_dataloader import NumpyLoader

logger = logging.getLogger(__name__)


class Trainer:
    def __init__(self, cfg):
        self.cfg = cfg
        self.logdir = cfg.logdir
        self.ckp_name = cfg.ckp_name

        self.model = hydra.utils.instantiate(cfg.model)
        self.training_data = hydra.utils.instantiate(cfg.train_dataset)
        self.data_loader = NumpyLoader(
            self.training_data,
            batch_size=cfg.train.batch_size,
            shuffle=True,
            num_workers=cfg.train.num_workers,
            pin_memory=False,
        )
        self.num_epochs = cfg.train.epochs
        self.save_model_freq = cfg.train.save_model_freq

    def run(self):
        writer = SummaryWriter(log_dir=self.logdir)

        for epoch in trange(self.num_epochs, desc=f"Train {self.ckp_name}"):

            cnt_batch = 0
            avg_info = {}

            for batch in self.data_loader:
                info = self.model.update(batch)

                for k, v in info.items():
                    if k in avg_info:
                        avg_info[k].append(v)
                    else:
                        avg_info[k] = [v]
                cnt_batch += 1

            for k, v in avg_info.items():
                writer.add_scalar(f"train/{k}", np.mean(v), epoch)

            if (1 + epoch) % 10 == 0:
                logger.info(f"Epoch{epoch+1}: loss = {np.mean(avg_info['loss'])}")

            if (1 + epoch) % self.save_model_freq == 0:

                self.model.save_ckpt(file_name=self.ckp_name + ".ckp")

        self.model.save_ckpt(file_name=self.ckp_name + ".ckp")


if __name__ == "__main__":

    OmegaConf.register_new_resolver("eval", eval, replace=True)

    # use line-buffering for both stdout and stderr
    sys.stdout = open(sys.stdout.fileno(), mode="w", buffering=1)
    sys.stderr = open(sys.stderr.fileno(), mode="w", buffering=1)

    @hydra.main(
        version_base=None,
        config_path=os.path.join(os.getcwd(), "cfg"),  # possibly overwritten by --config-path
    )
    def main(cfg: OmegaConf):
        OmegaConf.resolve(cfg)
        os.environ["CUDA_VISIBLE_DEVICES"] = cfg.devices
        app = Trainer(cfg)  # use other configs to create class object
        app.run()

    main()

    """
    python train.py --config-name=train_mlp_ik
    python train.py --config-name=train_diffusion_partial_ik
    """
